import multiprocessing
import random
from pprint import pprint
from mpi4py import MPI
import pandas as pd
import os
from expDesign import getCapacityOne, getCapacityTwo, getCapacityThree
#  一次一因子
# agent 属性：休息时间 工作时间 学习方式 性别 勤奋程度
# 干预策略：intervention_strategy， 比例 调控类型
# 交互方式：交互网络格式 概率 BA边 WS边 RG边
agents = [{
    'sleep_time': 200,
    'work_time': 2,
    'learning_type': 0,
    'sex': 0,
    'ati': 10
},
    {
        'sleep_time': 200,
        'work_time': 2,
        'learning_type': 0,
        'sex': 0,
        'ati': 10
    }
]
init_state = {
    'height': 0,
    'leaf_num': 0
}
params = {
    'test1': ['low', 'mid', '1'],
    'test2': ['low', 'mid', '1'],
    'test3': ['low', 'mid', '1']
}


def get_agent_percentage(agents):
    params['agent type'] = [i for i in range(len(agents))]
    params['agent type'].append(len(agents))


# seq为试验考虑的因子id（从1开始），param为参数
def getCapacity(seq, param):
    seq = seq - 1
    assert seq < len(param), '输入的次序应小于参数数量'
    assert seq >= 0, '次序应当不小于1'
    keys = []
    for key in param.keys():
        keys.append(key)

    # print(keys)
    capacity = {}
    for value in param[keys[seq]]:
        if keys[seq] not in capacity:
            capacity[keys[seq]] = []
        capacity[keys[seq]].append(value)
    for i in range(len(capacity[keys[seq]])):
        for item in param.items():
            if item[0] == keys[seq]:
                continue
            if item[0] not in capacity:
                capacity[item[0]] = []
            capacity[item[0]].append(item[1][0])
    return capacity


def run():
    get_agent_percentage(agents)
    # print(params['agent type'])
    capacity = getCapacity(seq=1, param=params)

    # print(capacity)


def transform_dict_to_3d_array(data_dict):
    result = []

    # 获取字典的第一个键对应的值，以确定需要重组的层数
    first_key_values = data_dict[next(iter(data_dict))]
    num_elements = len(first_key_values)

    # 遍历每个索引位置，构建新的数组
    for i in range(num_elements):
        new_array = {'first': '', 'second': '', 'third': ''}
        new_array['first'] = data_dict['test1'][i]
        new_array['second'] = data_dict['test2'][i]
        new_array['third'] = data_dict['test3'][i]
        result.append(new_array)

    return result


def exp_to_params(agents_id, num, strategies):
    # num为agent数量 agents_id为设计的方法， nums为一共提供的agent类型
    params = []
    for index, agent_id in enumerate(agents_id):
        type_ans = [int(agent_id) for _ in range(num)]
        ans = [
                MPI.COMM_WORLD,
                None,
                num,
                [200 for _ in range(num)],
                [5 for _ in range(num)],
                strategies[index],
                0.5,
                type_ans,
                [1 for _ in range(num)],
                [10 for _ in range(num)]
               ]
        params.append(ans)
    len_num = int(num / len(agents_id))
    type_ans = [int(n) for n in agents_id for _ in range(len_num)]
    while len(type_ans) < num:
        type_ans.append(random.choice(agents_id))
    ans = [
        MPI.COMM_WORLD,
        None,
        num,
        [200 for _ in range(num)],
        [5 for _ in range(num)],
        3,
        0.5,
        type_ans,
        [1 for _ in range(num)],
        [10 for _ in range(num)]
    ]
    params.append(ans)
    return params


def extract_system_utility_data(parent_folder_path, name_csv, data_column):
    """
    提取指定父文件夹下所有数字命名子文件夹中的指定CSV文件的指定列数据，
    并以字典形式返回，其中包含values和names两个键。

    参数:
    parent_folder_path (str): 包含数字命名子文件夹的父文件夹路径。
    name_csv (str): CSV文件的名称。
    data_column (str): 需要提取数据的列名。

    返回:
    dict: 包含"values"和"names"两个键，分别对应数据数组和实验编号名称列表。
    """
    system_utility_dict = {"value": [], "names": []}

    # 遍历父文件夹下的所有子文件夹
    for entry in os.listdir(parent_folder_path):
        sub_folder_path = os.path.join(parent_folder_path, entry)
        # 确保是文件夹且文件夹名称为数字
        if os.path.isdir(sub_folder_path) and entry.isdigit():
            # 构建CSV文件的完整路径
            csv_file_path = os.path.join(sub_folder_path, name_csv)
            # 检查CSV文件是否存在
            if os.path.isfile(csv_file_path):
                # 读取CSV文件
                df = pd.read_csv(csv_file_path)
                # 确保列名正确，然后提取指定列的数据
                if data_column in df.columns:
                    # 格式化键名称为“实验X”
                    experiment_name = f"实验{entry}"
                    # 将数据添加到values列表中
                    system_utility_dict["value"].append(df[data_column].tolist())
                    # 将实验名称添加到names列表中
                    system_utility_dict["names"].append(experiment_name)

    return system_utility_dict

if __name__ == "__main__":
    # multiprocessing.freeze_support()
    # run()
    # ans = exp_to_params([1, 2, 3], 10, [1, 1, 1, 2])
    # print(ans)
    ans = extract_system_utility_data('./info_log', 'utility.csv', 'system_utility')
    ans2 = getCapacityOne(1, params)
    ans3 = getCapacityThree(param=params, space=1)
    ans3 = transform_dict_to_3d_array(ans3)
    ans4 = transform_dict_to_3d_array(getCapacityTwo(3, 3, params))
    # print('ans4', ans4)
    # print('ans3', ans3)
    # print(ans2)
    # print(ans)
